{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import spacy\n",
    "import torch\n",
    "from transformers import pipeline\n",
    "\n",
    "from selfcheckgpt.modeling_selfcheck import SelfCheckLLMPrompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x14dedf3998f0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(28)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:07<00:00,  3.73s/it]\n"
     ]
    }
   ],
   "source": [
    "# We use Phi-2 2.7B SLM for inferencing\n",
    "pipe = pipeline(\"text-generation\", model=\"microsoft/phi-2\", device_map=\"auto\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "Give me the professional journey of Ashish Vaswani in detail.\n",
    "Answer:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'generated_text': 'Ashish Vaswani started his career as a software engineer in a small startup company. He worked hard and gained experience in various projects, which helped him climb up the corporate ladder. He then joined a multinational company as a project manager and was responsible for managing a team of developers. His leadership skills and dedication towards his work were recognized, and he was promoted to the position of a senior project manager.\\n\\nAfter a few years, Ashish decided to pursue his passion for entrepreneurship and started his own software development company. He faced many challenges in the initial stages, but with his determination and hard work, he was able to establish a'}]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# As per the original paper the response is generated with greedy decoding\n",
    "Response = pipe(prompt, do_sample=False, max_new_tokens=128, return_full_text=False)\n",
    "Response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'generated_text': \"Ashish Vaswani has always been passionate about the Indian film industry, and he pursued his interest by working as a producer and writer. He started his career by serving his family's business and then moved on to other jobs before finally entering the film industry.\\nAs a producer and screenwriter, Ashish has a diverse background, having worked on multiple projects such as music videos, TV commercials, short films, documentaries, and web series. He gained recognition for his work and has worked extensively in the Telugu film industry since 2017.\\nAshish's journey in the film industry has been a learning curve, and he has been involved\"}]\n"
     ]
    }
   ],
   "source": [
    "# The samples are generated for the same prompt with temperature as 1.\n",
    "N = 20\n",
    "Samples = pipe(\n",
    "    [prompt] * N,\n",
    "    temperature=1.0,\n",
    "    do_sample=True,\n",
    "    max_new_tokens=128,\n",
    "    return_full_text=False,\n",
    ")\n",
    "print(Samples[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "Response = Response[0][\"generated_text\"]\n",
    "Samples = [sample[0][\"generated_text\"] for sample in Samples]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mistral 7B became a gated repository so huggingface login is required to access it.\n",
    "from huggingface_hub import login\n",
    "\n",
    "HUGGINGFACE_TOKEN = \"...\"\n",
    "login(token=HUGGINGFACE_TOKEN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:08<00:00,  2.99s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SelfCheck-LLMPrompt (mistralai/Mistral-7B-Instruct-v0.2) initialized to device cuda\n"
     ]
    }
   ],
   "source": [
    "# We use Mistral 7B LLM to detect whether the response generated with Phi-2 LM is hallucinated or not using LLM Promting technique.\n",
    "llm_model = \"mistralai/Mistral-7B-Instruct-v0.2\"\n",
    "selfcheck_prompt = SelfCheckLLMPrompt(llm_model, device)\n",
    "\n",
    "# Option2: API access (currently only support client_type=\"openai\")\n",
    "# from selfcheckgpt.modeling_selfcheck_apiprompt import SelfCheckAPIPrompt\n",
    "# selfcheck_prompt = SelfCheckAPIPrompt(client_type=\"openai\", model=\"gpt-3.5-turbo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Ashish Vaswani started his career as a software engineer in a small startup company.', 'He worked hard and gained experience in various projects, which helped him climb up the corporate ladder.', 'He then joined a multinational company as a project manager and was responsible for managing a team of developers.', 'His leadership skills and dedication towards his work were recognized, and he was promoted to the position of a senior project manager.', 'After a few years, Ashish decided to pursue his passion for entrepreneurship and started his own software development company.', 'He faced many challenges in the initial stages, but with his determination and hard work, he was able to establish a']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:58<00:00,  9.79s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.   0.3  0.75 0.6  0.65 0.05]\n",
      "Hallucination Score: 0.5583333333333332\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "nlp = spacy.load(\"en_core_web_sm\")\n",
    "sentences = [\n",
    "    sent.text.strip() for sent in nlp(Response).sents\n",
    "]  # spacy sentence tokenization\n",
    "print(sentences)\n",
    "\n",
    "sent_scores_prompt = selfcheck_prompt.predict(\n",
    "    sentences=sentences,  # list of sentences\n",
    "    sampled_passages=Samples,  # list of sampled passages\n",
    "    verbose=True,  # whether to show a progress bar\n",
    ")\n",
    "\n",
    "print(sent_scores_prompt)\n",
    "print(\"Hallucination Score:\", np.mean(sent_scores_prompt))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
